import bisect
import math
import sys

import pylab
import matplotlib
from matplotlib import pyplot
from matplotlib.colors import LogNorm
import numpy

from h1d import Histogram1D

class Histogram2D(object):
    _rcparams = {'axes.grid': True,
                 'axes.labelsize': 'x-large',
                 'axes.titlesize': 'x-large',
                 'axes.formatter.limits': (-3, 3),                                 
                 'font.size': 24.0,
                 'figure.figsize': (8, 6),
                 'grid.color': 'w',
                 'text.latex.preamble': [r'\usepackage{amsmath}'],
                 'text.usetex': True,
                 'xtick.labelsize': 'large',
                 'xtick.direction': 'out',
                 'ytick.labelsize': 'large',
                 'ytick.direction': 'out',
                }

    _figprops = {'axes': (0.185, 0.165, 0.8, 0.72),
                 'entries': (1.25, -0.2),
                 'entries.fontsize': 20.0,
                }


    """
    Two-dimensional histogram object.
    
    Parameters
    ----------
    name   : string
        A name for the histogram
    nbins  : sequence of length 2
        (``number of bins in x``, ``number of bins in y``)
    arange : a sequence of two sequences of length 2
        ((``x axis minimum``, ``x axis maximum``), (``y axis minimum``, ``y axis maximum``))
    xlabel : string
        Label for the x-axis
    ylabel : string
        Label for the y-axis
    xunits : string 
        Units for the x-axis
    yunits : string 
        Units for the y-axis

    Optional Parameters
    -------------------
    title : string 
        Title of the histogram (default is "`xlabel` vs `ylabel`")
    fontsize  :  float | {'xx-small' , 'x-small' , 'small' , 'medium' , 'large' , 'x-large' , 'xx-large'}
        Size of text when drawn
    ext       : {'png' , 'gif' , 'jpg'}
        Default file extension to add to name when histogram is drawn and saved to file
    bufsize   : int
            length of internal data buffer; larger bufsizes are faster for large datasets but take up more memory
    assoc_cols     : 2-tuple of string
        Two column names (x and y) which are associated with this histogram
    selection : string
        A PyTables-style selection which is associated with this histogram
    logz      : bool
        Set the histogram to have a logarithmic z-axis when drawn            
    use_latex : True | False
        Turn on/off Matplotlib feeding image text through LaTeX

    """

    def __init__(self, name, nbins, aranges, xlabel, ylabel, xunits, yunits,
                 title='', fontsize='x-large', ext=['png'], 
                 bufsize=1000000, assoc_cols=[], logz=False, 
                 bins=None, use_latex=None):

        self.name      = name
        self.nxbins    = nbins[0]
        self.nybins    = nbins[1]
        self.xrange    = aranges[0]
        self.yrange    = aranges[1]
        self.xlabel    = xlabel
        self.ylabel    = ylabel
        self.xunits    = xunits
        self.yunits    = yunits
        self.title     = title
        self.tag       = ''
        self.fontsize  = fontsize
        self.ext       = ext
        self.assoc_cols = assoc_cols

        self.logz         = logz

        if use_latex is not None:
            self.use_latex(use_latex)

        self._i        = 0;
        self._tot_i    = 0;

        self._bufsize  = bufsize

        self.low_edges_x = None
        self.low_edges_y = None
        self.bin_width = (float(self.xrange[1] - self.xrange[0]) / self.nxbins,
                          float(self.yrange[1] - self.yrange[0]) / self.nybins)


        if bins is not None:
            self._bins  = bins
            self.low_edges_x, bin_width_x = numpy.linspace(self.xrange[0], 
                                                           self.xrange[1], 
                                                           self.nxbins, 
                                                           retstep=True)
            self.low_edges_y, bin_width_y = numpy.linspace(self.yrange[0], 
                                                           self.yrange[1], 
                                                           self.nybins, 
                                                           retstep=True)

            self._mean  = self._mean_from_bins()
            self._var   = self._var_from_bins()
            self._xmin, self._xmax = self._min_from_bins()
            self._ymin, self._ymax = self._max_from_bins()
            self._tot_i = numpy.sum(bins)

        else:
            self._bins  = None 
            self._mean   = 0.0 
            self._var    = 0.0
            self._xmin    = 1e250
            self._xmax    = -1e250
            self._ymin    = 1e250
            self._ymax    = -1e250

        self._buf  = None
        self._wbuf = None 

    def _setup_arrays(self):
        if self._bins is None:
            self._bins  = numpy.zeros((self.nxbins, self.nybins))

        if self._buf is None:
            self._buf   = numpy.zeros((self._bufsize, 2))

        if self._wbuf is None:
            self._wbuf  = numpy.ones((self._bufsize,))

        if self.low_edges_x is None:
            self.low_edges_x, bin_width_x = numpy.linspace(self.xrange[0], self.xrange[1], self.nxbins, retstep=True)

        if self.low_edges_y is None:
            self.low_edges_y, bin_width_y = numpy.linspace(self.yrange[0], self.yrange[1], self.nybins, retstep=True)

        
        
    def draw(self, hold=False, same=False, cmap='hot', **kwargs):
        """
        Draw this histogram.

        Optional Parameters
        -------------------
        cmap      : {'autumn', 'bone', 'cool', 'copper', 'flag', 'gray', 'hot', 'hsv', 'jet', 'pink', 'prism', 'spring', 'summer', 'winter', 'spectral'}
            Use colormap ``cmap``
        ext       : {'png' , 'gif' , 'jpg'}
            File extension to add to name when histogram is drawn and saved to file
        hold       : bool
            Draw histogram but do not write to file yet            
        logz      : bool
            Set the histogram to have a logarithmic z-axis when drawn                    
        name   : string
            A name for the histogram
        same       : bool
            Draw histogram atop another previously-held histogram
        xlabel : string
            Label for the x-axis
        xunits : string 
            Units for the x-axis
        ylabel : string
            Label for the y-axis
        yunits : string 
            Units for the y-axis
        """

        if self._bins is None:
            raise ValueError("Attempt to draw empty histogram `{0}'".format(self.name))
        
        self.flush()

        if self._tot_i == 0:
            raise ValueError("Attempt to draw empty histogram `{0}'".format(self.name))

        orig_attrs  = dict()
        local_attrs = list()
        for key in kwargs:
            try:
                orig_attrs[key] = getattr(self, key)
            except AttributeError:            
                local_attrs.append(key)

            setattr(self, key, kwargs[key])

        orig_rcparams = {}
        if not same:
            orig_rcparams.update(matplotlib.rcParams)
            matplotlib.rcParams.update(Histogram2D._rcparams)

            fig = pyplot.figure()
            axes = pyplot.axes(Histogram2D._figprops['axes'])
           

        X, Y = numpy.meshgrid(self.low_edges_x, self.low_edges_y)
        if self.logz:
            bins = self._bins.T + 1
            pyplot.pcolormesh(X, Y, bins, norm=LogNorm(vmin=bins.min(), vmax=bins.max()))
        else:
            pyplot.pcolormesh(X, Y, self._bins.T)
        
        if not same:
            pyplot.colorbar()
                    
            if not self.xunits:
                fmtstr = '{0}'
            else:
                fmtstr = '{0} ({1})'

            pyplot.xlabel(fmtstr.format(self.xlabel, self.xunits))

            if not self.yunits:
                fmtstr = '{0}'
            else:
                fmtstr = '{0} ({1})'

            pyplot.ylabel(fmtstr.format(self.ylabel, self.yunits))
        
            if not self.title:
                self.title = ' vs. '.join([self.xlabel, self.ylabel])
            
            if self.tag:
                self.title += ' (' + self.tag + ')' 

            pyplot.title(self.title)
                
            axes = pyplot.gca()
            pyplot.text(Histogram2D._figprops['entries'][0], 
                        Histogram2D._figprops['entries'][1], 
                        'Entries: {0:.3g}'.format(self.nentries()),
                        transform=axes.transAxes,
#                        bbox=dict(boxstyle='square', facecolor='white'),
                        horizontalalignment='right',
                        fontsize=Histogram2D._figprops['entries.fontsize'])

            pyplot.axis(xmin=self.xrange[0], xmax=self.xrange[1], ymin=self.yrange[0], ymax=self.yrange[1])


        try:
            pyplot.set_cmap(getattr(pylab.cm, cmap))
#            getattr(pyplot, cmap)()
        except AttributeError:
            sys.stderr.write("Histogram2D.draw() failed: no colormap `{0}' in pyplot\n".format(cmap))
            return

        pyplot.grid(b='on')

#        pyplot.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
        retname = None
        if not hold:
            if self.name[-3:].lower() in 'png jpg gif pdf svg eps'.split():
                pyplot.savefig(self.name)
                retname = self.name
            else:
                retname = []
                for ext in self.ext:
                    outname = '.'.join([self.name, ext])
                    pyplot.savefig(outname)
                    retname.append(outname)

            pyplot.clf()

        for key in orig_attrs:
            setattr(self, key, orig_attrs[key])
        
        for attr in local_attrs:
            delattr(self, attr)

        return retname
    
    def save_and_clear(self, **kwargs):
        """
        Save the current contents of the canvas to file and clear the canvas.

        Optional parameters
        -------------------
        ext -- {'png', 'gif', 'jpg'}
            Extension of the output image
        name -- string
            Name of the output file

        """
        orig_attrs  = dict()
        local_attrs = list()
        for key in kwargs:
            try:
                orig_attrs[key] = getattr(self, key)
            except AttributeError:
                local_attrs.append(key)

            setattr(self, key, kwargs[key])


        retname = None

#        pyplot.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)

        if self.name[-3:].lower() in 'png jpg gif pdf svg'.split():
            pyplot.savefig(self.name)
            retname = self.name
        else:
            retname = []
            for ext in self.ext:
                outname = '.'.join([self.name,ext])
                pyplot.savefig(outname)
                retname.append(outname)

        for key in orig_attrs:
            setattr(self, key, orig_attrs[key])
        
        for attr in local_attrs:
            delattr(self, attr)

        pyplot.clf()

        return retname


    def fill(self, value, weight=1, loud=False):
        """
        Fill a histogram with values.
                
        Parameters
        ----------
        
        value : ``(x, y)`` | sequence of ``(x, y)``
            Values to be binned

        Optional Parameters
        -------------------
        weight : ``(x_weight, y_weight)`` | sequence of ``(x_weight, y_weight)``
            Weights to apply when binning values
        loud   : bool
            Print a progress indicator when binning multiple values
        """

        if self._buf is None:
            self._setup_arrays()

        try:
            self._buf[self._i] = value
        except IndexError:
            self.flush()
            try:
                self._buf[self._i] = value
            except ValueError:
                return self._fill_array(numpy.asarray(value), weight, loud)
        except ValueError:
            return self._fill_array(numpy.asarray(value), weight, loud)

        self._wbuf[self._i] = weight
        self._i += 1

    def _fill_array(self, values, weights, loud):                       
        if not len(values):
            return

        valshape = numpy.shape(values)

        if len(valshape) != 2 or 2 not in valshape:
            msg = "Trying to fill 2D histogram with array of shape `{0}'; valid shapes are (2,n) or (n,2)"
            raise ValueError(msg.format(valshape))
        
        if valshape[1] != 2:
            values = values.T

        try:
            x = iter(weights)
        except TypeError:
            pass
        else:
            return self._fill_array_weights(values, weights)
        
        self.flush()
        
        val_len = len(values)

        if val_len < self._bufsize:
            self._i = val_len
            self._buf[:self._i] = values[:self._i]
            return

        for i in xrange(0, val_len, self._bufsize):
            if loud:
                sys.stderr.write("Histogram1D `{0}': {1}/{2} events binned\r".format(self.name, i, val_len))
                sys.stderr.flush()
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self.flush()

    def _fill_array_weights(self, values, weights):
        self.flush()
        
        val_len  = len(values)
        if val_len < self._bufsize:
            self._i = len(values)
            self._buf[:self._i]  = values
            self._wbuf[:self._i] = weights
            return

        for i in xrange(0, val_len, self._bufsize):
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self._wbuf[:self._i] = weights[i : i + self._bufsize]
            self.flush()
    

    def clear(self):
        """Set all bin contents to zero."""

        if self._buf is not None:
            self._buf   = numpy.zeros((self._bufsize,))

        if self._bins is not None:
            self._bins  = numpy.zeros((self.nxbins,self.nybins))

        self._mean   = 0.0 
        self._var    = 0.0
        self._xmin    = 1e250
        self._xmax    = -1e250
        self._ymin    = 1e250
        self._ymax    = -1e250
        self._i      = 0
        self._tot_i  = 0


    def flush(self):
        """
        Bin events in buffer and clear the buffer.
        """

        if self._buf is None or self._bins is None:
            self._setup_arrays()

        if self._i == 0:
            return
        
        buf_to_flush = self._buf[:self._i]

        bufxmax = max(buf_to_flush.T[0])
        if bufxmax > self._xmax:
            self._xmax = bufxmax
        bufxmin = min(buf_to_flush.T[0])
        if bufxmin < self._xmin:
            self._xmin = bufxmin

        bufymax = max(buf_to_flush.T[1])
        if bufymax > self._ymax:
            self._ymax = bufymax
        bufymin = min(buf_to_flush.T[1])
        if bufymin < self._ymin:
            self._ymin = bufymin

        
        self._mean = (self._mean * self._tot_i + numpy.mean(buf_to_flush) * self._i) / (self._tot_i + self._i)
        self._var  += sum((buf_to_flush - self._mean)**2)
 
        bufbins, edges = numpy.histogramdd(buf_to_flush,
                                           bins=(self.nxbins, self.nybins),
                                           weights=self._wbuf[:self._i],
                                           range=(self.xrange, self.yrange),
                                           )
        
        self._bins += bufbins
        self._tot_i += self._i
        self._i = 0
        
    def project_to_y(self, x_start=0, x_stop=None):
        """
        Project the histogram onto the x-axis

        Optional Parameters
        -------------------
        bin_start : int
            Bin index of lower bound of projection
        bin_stop  : int
            Bin index of upper bound of projection

        Returns
        -------
        a Histogram1D object.
        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()
        
        bin_start, iy = self.bin_index(x_start, 0)
  
        if x_stop is None:
            bin_stop = self._bins.shape[0]
        else:
            bin_stop, iy = self.bin_index(x_stop, 0)

        bins = numpy.sum(self._bins[bin_start:bin_stop], 0)
        return Histogram1D(self.name + "_py", len(bins), self.yrange, self.ylabel, self.yunits,
                           self.ylabel, bins=bins)

    def project_to_x(self, bin_start=0, bin_stop=None):
        """
        Project the histogram onto the x-axis

        Optional Parameters
        -------------------
        bin_start : int
            Bin index of lower bound of projection
        bin_stop  : int
            Bin index of upper bound of projection

        Returns
        -------
        a Histogram1D object.

        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        if bin_stop is None:
            bin_stop = self._bins.shape[1]

        bins = numpy.sum(self._bins.T[bin_start:bin_stop], 0)
        return Histogram1D(self.name + "_px", len(bins), self.xrange, self.xlabel, self.xunits,
                           self.xlabel, bins=bins)

    
    def set_buffer_size(self, bufsize):
        """
        Set internal buffer size to new value.
        
        Parameters
        ----------
        
        bufsize : int
            desired new buffer size
        """

        self._bufsize = bufsize

        if self._buf is not None:            
            self.flush()
            
            self._buf     = numpy.zeros((self._bufsize,))
            self._wbuf    = numpy.ones((self._bufsize,))


    def __getstate__(self):
        self.flush()

        outdict = {}
        outdict.update(self.__dict__)
        
        outdict['low_edges_x'] = None
        outdict['low_edges_y'] = None
        outdict['_bins'] = None
        outdict['_buf'] = None
        outdict['_wbuf'] = None
        
        return outdict

    def _min_from_bins(self):
        return (0,0)

    def _max_from_bins(self):
        return (0,0)

    def _mean_from_bins(self):
    #   1D case
    #   mean = sum([self._bins[i]*(self.low_edges[i] + (self.bin_width/2.0)) for i in xrange(self.nbins)]) / sum(self._bins)
        return 0.0

    def _var_from_bins(self):
    #   1D case
    #   var = sum([self._bins[i]*(self.low_edges[i] + (self.bin_width/2.0) - self._mean)**2 for i in xrange(self.nbins)])
        return 0.0


    def bins(self):
        """Returns a reference to the array of bin contents."""

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        return self._bins

    def max_bin_index(self):
        """
        Returns the indices ``(ix, iy)`` of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()

        flat_index = numpy.argmax(self._bins)

        return numpy.unravel_index(numpy.argmax(self._bins), self._bins.shape)                


    def max_bin_content(self):
        """
        Returns the number of counts in the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        i, j = self.max_bin_index()
        return self._bins[i][j]

    def max_bin_xy(self):
        """
        Returns the x-y coordinates ``(x, y)`` of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        return self.bin_center(*self.max_bin_index())

    def bin_center(self, ix, iy):
        """
        Returns the x-y coordinates ``(x, y)`` of the center of the bin at ``(ix, iy)``.
        """

        if self.low_edges_x is None or self.low_edges_y is None:
            self._setup_arrays()

        return (self.low_edges_x[ix] + self.bin_width[0]*0.5, self.low_edges_y[iy] + self.bin_width[1]*0.5)

    def bin_index(self, x, y):
        """
        Return the indices ``(ix, iy)`` of the bin that ``(x, y)`` would be placed into.
        """

        if self.low_edges_x is None or self.low_edges_y is None:
            self._setup_arrays()
        
        return (bisect.bisect(self.low_edges_x, x), bisect.bisect(self.low_edges_y, y))
    
    def data_xmax(self):
        """
        Returns the current maximum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmax

    def data_xmin(self):
        """
        Returns the current minimum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmin

    def data_ymax(self):
        """
        Returns the current maximum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymax

    def data_ymin(self):
        """
        Returns the current minimum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymin

    def data_mean(self):
        """
        Returns the mean ``(mu_x, mu_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._mean

    def data_var(self):
        """
        Returns the variance ``(var_x, var_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._var

    def data_std(self):
        """
        Returns the standard deviation ``(std_x, std_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return math.sqrt(self._var / float(self._tot_i))

    def nentries(self):
        """
        Returns the total number of values binned.
        """

        if self._bins is None:
            return 0
        
        self.flush()
        return self._tot_i

        counts = numpy.sum(self._bins)

        if counts < int(counts) > 0.0001:
            return counts
        else:
            return int(counts)


    def __len__(self):
        return self.nentries()

                    
    def __iadd__(self, rhs):     
        if self._bins is None:
            self._setup_arrays()

        if rhs._bins is None:
            rhs._setup_arrays()

        self.flush()
        rhs.flush()

        if len(rhs) == 0:
            return self
        
        tot_entries = self._tot_i + rhs._tot_i
        self._mean = self._mean * (self._tot_i / float(tot_entries)) + rhs._mean * (rhs._tot_i / float(tot_entries))
        self._var  = self._var * (self._tot_i / float(tot_entries)) + rhs._var * (rhs._tot_i / float(tot_entries))

        self._tot_i = tot_entries
        
        if self._xmax < rhs._xmax:
            self._xmax = rhs._xmax
        
        if self._xmin > rhs._xmin:
            self._xmin = rhs._xmin        

        if self._ymax < rhs._ymax:
            self._ymax = rhs._ymax
        
        if self._ymin > rhs._ymin:
            self._ymin = rhs._ymin        

        self._bins += rhs._bins
        
        return self

    def use_latex(self, value):
        from matplotlib import rcParams
        rcParams['text.usetex'] = value


    def __str__(self):
        s = "Histogram2D `{0}'; {1} bins from {2} to {3} {4} in X; {5} bins from {6} to {7} {8} in Y; {9} entries"
        return s.format(self.name, self.nxbins, self.xrange[0], self.yrange[1], self.xunits, 
                        self.nybins, self.yrange[0], self.yrange[1], self.yunits,
                        self.nentries())
